Add use_causal_mask to MultiHeadAttention and approximate GELU - #466
Closed
Dobiasd wants to merge 1 commit into
Closed
Add use_causal_mask to MultiHeadAttention and approximate GELU#466Dobiasd wants to merge 1 commit into
Dobiasd wants to merge 1 commit into
Conversation
The MultiHeadAttention layer now respects the per-call use_causal_mask=True kwarg: scores at (t, k) for k > t are set to -inf before softmax. The converter extracts the flag from the layer's first inbound node and bakes it into the JSON config so the runtime can pick it up. GELU gains the tanh-approximation form. A serializable gelu_approximate helper in convert_model lets Keras models reference it via activation=gelu_approximate; the converter rewrites that to a plain "gelu" activation with approximate=True in the config, which the C++ gelu_layer applies via the appropriate formula. Both features are exercised by new cases in the exhaustive test model. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Owner
Author
|
Closing — landed as an experiment, not pursuing for merge. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
MultiHeadAttentionnow respectsuse_causal_mask=True. The converter extracts the flag from the layer's first inbound node (it's a per-call kwarg, not stored inget_config()), bakes it into the JSON, and the C++ runtime applies a triangular mask on the score tensor before softmax.gelugains anapproximate=Truepath (tanh form). A serializablegelu_approximatehelper inconvert_modellets Keras models reference it viaactivation=gelu_approximate; the converter rewrites the activation to plain"gelu"withapproximate=Truein the layer config, and the runtime picks the right formula. Existing exact-GELU users are unaffected (defaultapproximate=False).test_model_exhaustive(three new MHA cases withuse_causal_mask=True, two newDensecases for exact / approximate GELU).Together these are the runtime building blocks needed to load decoder-only LLMs (e.g. GPT-2) into frugally-deep. The actual GPT-2 example tooling is split out into a follow-up PR to keep this one focused.
Test plan
test_model_exhaustivepasses (with new MHA causal-mask + GELU approximate cases)test_model_embedding,test_model_recurrent,test_model_sequential,test_model_variable,test_model_autoencoderall still pass after the converter changes (they invokerewrite_custom_activationsandinject_mha_call_kwargseven when no rewrites apply).🤖 Generated with Claude Code